#
# peerscape.lib.cipher
#
# Copyright 2008 Helsinki Institute for Information Technology
# and the authors.
#
# Authors: Ken Rimey <rimey@hiit.fi>
#

# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

"""
Easy-to-use AES encryption

>>> cipher = Cipher('my password')
>>> plaintext = 'This is my secret message.'
>>> ciphertext = cipher.encrypt(plaintext)
>>> cipher.decrypt(ciphertext) == plaintext
True
>>> len(plaintext), len(ciphertext)
(26, 48)
>>> cipher.encrypt(plaintext) != ciphertext
True
"""

from hashlib import md5
from os import urandom

try:
    from Crypto.Cipher import AES
    def block_cipher_fast(key, IV):
        return AES.new(key, AES.MODE_CBC, IV)
except ImportError:
    block_cipher_fast = None

try:
    from rijndael import rijndael
    def block_cipher_slow(key, IV):
        return CBC(rijndael(key), IV)
except ImportError:
    block_cipher_slow = None

default_block_cipher = block_cipher_fast or block_cipher_slow

class Cipher(object):
    def __init__(self, key, IV=None, block_cipher=None):
        self.block_cipher = block_cipher or default_block_cipher

        if len(key) != 16:
            key = md5(key).digest()

        if IV is not None and len(IV) != 16:
            IV = md5(IV).digest()

        self.key = key
        self.IV = IV

    def encrypt(self, plaintext):
        padded_plaintext = pad_string(plaintext)
        if self.IV is not None:
            cipher = self.block_cipher(self.key, self.IV)
            return cipher.encrypt(padded_plaintext)
        else:
            IV = urandom(16)
            cipher = self.block_cipher(self.key, IV)
            return IV + cipher.encrypt(padded_plaintext)

    def decrypt(self, ciphertext):
        if self.IV is not None:
            cipher = self.block_cipher(self.key, self.IV)
            padded_plaintext = cipher.decrypt(ciphertext)
        else:
            IV, ciphertext = ciphertext[:16], ciphertext[16:]
            cipher = self.block_cipher(self.key, IV)
            padded_plaintext = cipher.decrypt(ciphertext)
        return unpad_string(padded_plaintext)

class CBC(object):
    def __init__(self, cipher, IV):
        self.cipher = cipher
        self.IV = IV

    def encrypt(self, s):
        return map16(self.encrypt16, s)

    def decrypt(self, s):
        return map16(self.decrypt16, s)

    def encrypt16(self, plaintext):
        ciphertext = self.cipher.encrypt(xor(self.IV, plaintext))
        self.IV = ciphertext
        return ciphertext

    def decrypt16(self, ciphertext):
        plaintext = xor(self.IV, self.cipher.decrypt(ciphertext))
        self.IV = ciphertext
        return plaintext

def _test(plaintext, key, IV=None, block_cipher=None):
    """
    >>> block_cipher_fast is not None  # Need pycrypto.
    True
    >>> block_cipher_slow is not None
    True

    >>> plaintext = urandom(1024)
    >>> ciphertext1 = _test(plaintext, 'my password', 'deterministic',
    ...                     block_cipher=block_cipher_fast)
    >>> ciphertext2 = _test(plaintext, 'my password', 'deterministic',
    ...                     block_cipher=block_cipher_slow)
    >>> ciphertext1 == ciphertext2
    True
    >>> len(ciphertext1) - 1024
    16
    """
    cipher = Cipher(key, IV, block_cipher)
    ciphertext = cipher.encrypt(plaintext)
    assert cipher.decrypt(ciphertext) == plaintext
    return ciphertext

def map16(f, s):
    """
    >>> def capitalize(block): return block.capitalize()
    >>> map16(capitalize, 'this is a nice test sentence with 48 characters.')
    'This is a nice tEst sentence witH 48 characters.'
    """
    results = []
    for i in range(0, len(s), 16):
        block = s[i : i + 16]
        results.append(f(block))
    return ''.join(results)

def xor(s1, s2):
    """
    >>> xor(xor('foo', 'bar'), 'bar')
    'foo'
    """
    assert len(s1) == len(s2)
    result = [chr(ord(a) ^ ord(b)) for a, b in zip(s1, s2)]
    return ''.join(result)

def pad_string(s):
    n = 16 - len(s) % 16        # 1..16
    return s + chr(n) * n

def unpad_string(s):
    if s:
        n = ord(s[-1])
        if 1 <= n <= 16 and s[-n:] == chr(n) * n:
            return s[:-n]

    raise ValueError('Bad padding.')

def _test_padding(s):
    """
    >>> for n in range(50):
    ...     _test_padding('Z' * n)
    """
    s1 = pad_string(s)
    assert len(s1) % 16 == 0
    s2 = unpad_string(s1)
    assert s2 == s

if __name__ == '__main__':
    import doctest
    doctest.testmod()
